"""
LangChain 回调事件学习Demo - 单事件版本
专注于理解 on_llm_new_token 事件的工作原理
"""
import os

from dotenv import load_dotenv
from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema import HumanMessage
from langchain_openai import ChatOpenAI

# 加载环境变量
load_dotenv()


class LearningCallback(BaseCallbackHandler):
    """
    学习专用的回调处理器
    这个类只处理 on_llm_new_token 一个事件，便于理解核心概念

    继承自 BaseCallbackHandler 意味着我们只需要重写我们关心的事件方法
    其他事件可以完全忽略，LangChain 会自动处理
    """

    def on_llm_new_token(self, token: str, **kwargs):
        """
        🔥 核心方法：当语言模型生成每个新 token 时触发

        什么是 token？
        - 在 NLP 中，token 是文本的基本单位
        - 英文中可能是一个单词或标点："Hello"、"world"、"!" 都是 token
        - 中文中可能是一个字或词："人工"、"智能"、"。" 都是 token

        这个方法是如何工作的？
        1. LLM 开始生成文本时，不是一次性生成完整答案
        2. 而是一个 token 一个 token 地生成
        3. 每生成一个 token，这个方法就被调用一次
        4. 我们可以在这个时机做各种处理

        :param token: 新生成的 token（字或词）
        :param kwargs: 其他可选参数，包含执行上下文信息
        """

        # 🖨️ 最简单的处理：直接打印 token
        # end="" 表示打印后不换行，让 token 连续输出
        # flush=True 表示立即刷新输出缓冲区，实现"实时"效果
        print(token, end="", flush=True)


def demo_basic_learning():
    """
    基础学习演示：理解 on_llm_new_token 的基本工作原理
    """
    print("🎯 学习目标：理解 on_llm_new_token 事件")
    print("=" * 60)

    # 第一步：创建语言模型实例
    # 💡 注意：streaming=True 是触发 on_llm_new_token 的关键！
    llm = ChatOpenAI(api_key=os.getenv("MODELSCOPE_API_KEY"), openai_api_base=os.getenv("MODELSCOPE_API_BASE"), model="Qwen/Qwen2.5-VL-7B-Instruct", streaming=True, temperature=0.7)

    # 第二步：准备输入内容
    # HumanMessage 表示用户输入的消息
    messages = [
        HumanMessage(content="用简单的语言解释什么是机器学习？")
    ]

    print("📝 用户问题：用简单的语言解释什么是机器学习？")
    print("🔄 开始流式生成回答...")
    print("-" * 40)
    print("🤖 AI 回答: ", end="")  # end="" 让回答紧接着这个提示输出

    # 第三步：调用模型并传入我们的回调处理器
    response = llm.invoke(
        messages,  # 输入消息
        config={
            "callbacks": [LearningCallback()]  # 🔑 关键：传入回调处理器
        }
    )

    # 第四步：处理完成后的输出
    print("\n" + "-" * 40)
    print("✅ 生成完成！")
    print("💡 学习要点：")
    print("   - 你看到了文本是如何一个词一个词生成的")
    print("   - 这就是 on_llm_new_token 事件的威力")
    print("   - 每个 token 生成时，我们的回调方法都被调用一次")


def demo_advanced_learning():
    """
    进阶学习：在 on_llm_new_token 中添加更多功能
    """
    print("\n" + "🎯 进阶学习：增强的回调处理器")
    print("=" * 60)

    class AdvancedLearningCallback(BaseCallbackHandler):
        """带统计功能的进阶回调处理器"""

        def __init__(self):
            """
            初始化方法
            在这里设置计数器和其他状态变量
            """
            self.token_count = 0  # 统计生成的 token 数量
            self.start_time = None  # 记录开始时间

        def on_llm_start(self, serialized, prompts, **kwargs):
            """
            💡 额外事件：LLM 开始处理时触发
            虽然不是必须的，但可以帮我们更好地理解执行流程

            :param serialized: 模型的序列化信息
            :param prompts: 输入的提示词列表
            """
            import time
            self.start_time = time.time()  # 记录开始时间
            print(f"⏰ LLM 开始处理...")

        def on_llm_new_token(self, token: str, **kwargs):
            """
            增强的 token 处理方法
            现在除了输出，还进行统计
            """
            self.token_count += 1  # 每生成一个 token 就计数

            # 添加一些视觉效果，让输出更有趣
            if self.token_count % 5 == 0:  # 每5个token显示一次计数
                print(f" [{self.token_count}]", end="", flush=True)

            print(token, end="", flush=True)

        def on_llm_end(self, response, **kwargs):
            """
            💡 额外事件：LLM 处理完成时触发
            用于显示统计结果
            """
            import time
            end_time = time.time()
            duration = end_time - self.start_time

            print(f"\n" + "-" * 40)
            print(f"📊 生成统计:")
            print(f"   总 token 数: {self.token_count}")
            print(f"   总耗时: {duration:.2f} 秒")
            print(f"   平均速度: {self.token_count / duration:.1f} token/秒")

    # 使用进阶回调
    llm = ChatOpenAI(api_key=os.getenv("MODELSCOPE_API_KEY"), openai_api_base=os.getenv("MODELSCOPE_API_BASE"), model="Qwen/Qwen2.5-VL-7B-Instruct", streaming=True, temperature=0.7)

    messages = [HumanMessage(content="用三个关键词描述人工智能的未来")]

    print("📝 用户问题：用三个关键词描述人工智能的未来")
    print("🔄 开始生成（带统计功能）...")
    print("-" * 40)
    print("🤖 AI 回答: ", end="")

    # 创建回调实例，这样我们可以在执行后访问它的统计信息
    advanced_callback = AdvancedLearningCallback()

    response = llm.invoke(
        messages,
        config={"callbacks": [advanced_callback]}  # 传入回调实例
    )


def debug_demo():
    """
    调试演示：理解事件触发的时机和参数
    """
    print("\n" + "🎯 调试学习：查看事件详情")
    print("=" * 60)

    class DebugCallback(BaseCallbackHandler):
        """用于调试的回调处理器"""

        def on_llm_new_token(self, token: str, **kwargs):
            """
            调试版本的 token 处理
            显示每个 token 的详细信息
            """
            # 显示 token 的详细信息
            print(f"\n🔍 Token 详情:")
            print(f"   内容: '{token}'")
            print(f"   长度: {len(token)} 字符")
            print(f"   类型: {type(token)}")

            # 显示其他参数（如果有的话）
            if kwargs:
                print(f"   其他参数: {list(kwargs.keys())}")

            # 正常输出 token
            print(f"   输出: {token}", end="", flush=True)

    # 注意：这个演示会输出很多调试信息，可能看起来比较乱
    # 但它能帮你理解每个 token 的生成过程

    llm = ChatOpenAI(api_key=os.getenv("MODELSCOPE_API_KEY"), openai_api_base=os.getenv("MODELSCOPE_API_BASE"), model="Qwen/Qwen2.5-VL-7B-Instruct", streaming=True, temperature=0.7)

    messages = [HumanMessage(content="说'你好'")]

    print("📝 调试演示：观察每个 token 的生成")
    print("🔄 开始调试输出...")

    response = llm.invoke(
        messages,
        config={"callbacks": [DebugCallback()]}
    )


def main():
    """
    主函数：运行所有学习演示
    """
    print("🚀 LangChain 回调事件学习之旅")
    print("专注于理解 on_llm_new_token 事件")
    print()

    # 演示1：基础理解
    demo_basic_learning()

    # 演示2：进阶功能
    demo_advanced_learning()

    # 演示3：调试理解（可选，输出较多）
    # debug_demo()

    print("\n" + "🎉 学习总结")
    print("=" * 60)
    print("✅ 你学会了：")
    print("   - on_llm_new_token 事件在每个 token 生成时触发")
    print("   - streaming=True 是触发该事件的关键设置")
    print("   - 可以通过回调实现流式输出、统计、监控等功能")
    print("   - 回调处理器可以同时处理多个事件")
    print()
    print("💡 下一步建议：")
    print("   - 尝试修改回调，添加颜色、动画等效果")
    print("   - 学习其他回调事件（on_chain_start, on_tool_start 等）")
    print("   - 实践将回调用于日志记录或性能监控")


# 运行主程序
if __name__ == "__main__":
    main()
